import h5py
import pickle
from multiprocessing import Pool
import os
import scipy
from tqdm import tqdm
import mat73
import json
f = h5py.File('./tmp/a.h5','w')   #创建一个h5文件，文件指针是f  
f['ACE$ilugolyigiugiug'] = [1,2,3,4,5]
              #将数据写成data的键值
f.close()                        #关闭文件

f = h5py.File('./tmp/a.h5','r')   #打开h5文件  
# print f.keys()                 #查看所有的键
b = f['data'][:]                 #取出键名为data的键值
f.close()
print(b.shape)           # (10000, 5000)

import pickle

class Func4MultiProcessing: # 把多线程函数的逻辑部分打包进这个class，方便统一接口。
    def zuco2eeg(self, zuco2path, task_name, sbj_name, tn):
        file_name = f"{zuco2path}/{task_name}/Matlab files/results{sbj_name}_{tn}.mat"
        print(f"---==Reading[{file_name}]==---")
        text_ = mat73.loadmat(file_name)['sentenceData']['content']

        folder = f"{zuco2path}/{task_name}/Raw data/{sbj_name}/"
        filename_li = [f"{sbj_name}_{tn}{_}_EEG.mat" for _ in range(1, 8)]
        eeg_seg = []
        for file_name in tqdm(filename_li):
            print(f"---==Processing[{folder + file_name}]==---")
            data = scipy.io.loadmat(folder + file_name, squeeze_me=True, struct_as_record=False)
            raweeg = data["EEG"].data
            if "EEG" in data.keys():
                event = data["EEG"].event
            else:
                event = data["event"]


            cnt_10, cnt_11, cnt_12, cnt_13 = None, None, None, None
            cnt = 0
            for each in event:
                if each.value == "trigger":
                    if each.type.strip() == "10":
                        cnt_10 = each.latency
                    elif each.type.strip() == "11":
                        cnt_11 = each.latency
                        eeg_seg.append(raweeg[:, cnt_10:cnt_11+1])
                        cnt_10, cnt_11 = None, None
                        cnt += 1
                    elif each.type.strip() == "12":
                        cnt_12 = each.latency
                    elif each.type.strip() == "13":
                        cnt_13 = each.latency
                        eeg_seg.append(raweeg[:, cnt_12:cnt_13+1])
                        cnt_12, cnt_13 = None, None
                        cnt += 1
        assert len(eeg_seg) == len(text_)
        tmpdic = {"eeg_seg": eeg_seg, "text_": text_}
        pickle.dump(tmpdic, open(f"./tmp/zuco2_{task_name}_{sbj_name}.pkl", "wb"))
        print(f"---==Save [./tmp/zuco2_{task_name}_{sbj_name}.pkl]==---")

    @staticmethod
    def error_back(eb):
        print(f'error: {str(eb)}')


def mutipro_start(func, zuco2path, task_name, sbj_names_, tn):
    working_thread = []
    workernum = len(sbj_names_)
    with Pool(processes=workernum) as pool:
        for sbj_name in sbj_names_:
            w = pool.apply_async(
                func=func.zuco2eeg,
                args=(zuco2path, task_name, sbj_name, tn),
                error_callback=func.error_back
            )
            working_thread.append(w)
        for _ in range(workernum):
            working_thread[_].get()


def data_prepare_zuco2(zuco2path="./rawdata/zuco2", workernum=4):
    eeg, text = [], []
    for task_name in os.listdir(zuco2path):
        if "NR" in task_name:
            tn = "NR"
        elif "TSR" in task_name:
            tn = "TSR"
        else:
            continue
        sbj_names = list(os.listdir(f"{zuco2path}/{task_name}/Raw data"))
        for i in range(0, len(sbj_names), workernum):
            sbj_names_ = sbj_names[i:i + workernum]

            func = Func4MultiProcessing()
            mutipro_start(func, zuco2path, task_name, sbj_names_, tn)

    for tmpfile in os.listdir("./tmp"):
        if "zuco2" not in tmpfile:
            continue
        dic = pickle.load(open(f"./tmp/{tmpfile}", "rb"))
        eeg_seg, text_ = dic["eeg_seg"], dic["text_"]
        eeg += eeg_seg
        text += text_

    dic = {"eeg": eeg, "text": text}
    pickle.dump(dic, open("./data/zuco2.pkl", "wb"))
    print("---==Save [./data/zuco2.pkl]==---")